{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./sample_data/labeled_black/dns2tcp_cmd.pcap.txt\n",
      "./sample_data/labeled_black/dns2tcp_sendfile9.pcap.txt\n",
      "./sample_data/labeled_black/dns2tcp_when_use_ssh.txt\n",
      "./sample_data/labeled_black/dnscapy_scp.pcap.txt\n",
      "./sample_data/labeled_black/dnscapy_scp2.pcap.txt\n",
      "./sample_data/labeled_black/dnscat2_when_exec_command_rm_file.txt\n",
      "./sample_data/labeled_black/dnscat2_when_idle.txt\n",
      "./sample_data/labeled_black/download_dnscat2_file12.pcap.txt\n",
      "./sample_data/labeled_black/download_dnscat2_file13.pcap.txt\n",
      "./sample_data/labeled_black/download_dnscat2_file14.pcap.txt\n",
      "./sample_data/labeled_black/iodine_direct_scp17_base128.pcap.txt\n",
      "./sample_data/labeled_black/iodine_direct_scp3_base64.pcap.txt\n",
      "./sample_data/labeled_black/iodine_direct_ssh4_base32.pcap.txt\n",
      "./sample_data/labeled_black/iodine_direct_ssh6_base128.pcap.txt\n",
      "./sample_data/labeled_black/iodine_direct_ssh6_base64.pcap.txt\n",
      "./sample_data/labeled_black/iodine_direct_ssh6_base64u.pcap.txt\n",
      "./sample_data/labeled_black/iodine_direct_ssh9_base32_again.pcap.txt\n",
      "./sample_data/labeled_black/iodine_idle_direct_idle44.pcap.txt\n",
      "./sample_data/labeled_black/nbtoo_dnscat_file7.pcap.txt\n",
      "./sample_data/labeled_black/ozyman_idle.pcap.txt\n",
      "./sample_data/labeled_black/ozyman_idle2.pcap.txt\n",
      "./sample_data/labeled_black/ozyman_idle3.pcap.txt\n",
      "./sample_data/labeled_black/tcp-over-dns-idle.pcap.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-ctripgslb.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-dlgslb.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-gosuncdn.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-mccdnglb.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-mmycdn.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-ruisucdn.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-spcdntip.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-0-tcdnvod.com.txt\n",
      "./sample_data/labeled_cdn/2017-8-2-8-tcdnvod.com.txt\n",
      "./sample_data/labeled_white/2017-8-15-0-henanpeace.org.cn.txt\n",
      "./sample_data/labeled_white/2017-8-15-8-qichedaquan.com.txt\n",
      "./sample_data/labeled_white/2017-8-16-11-lse.ac.uk.txt\n",
      "./sample_data/labeled_white/2017-8-16-9-dicp.ac.cn.txt\n",
      "./sample_data/labeled_white/2017-8-2-0-365yg.com.txt\n",
      "./sample_data/labeled_white/2017-8-2-0-bilibiligame.net.txt\n",
      "X len: 4439 Y len: 4439\n",
      "<class 'numpy.ndarray'>\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding_14 (Embedding)     (None, 256, 64)           2752      \n",
      "_________________________________________________________________\n",
      "bidirectional_11 (Bidirectio (None, 32)                7776      \n",
      "_________________________________________________________________\n",
      "dense_19 (Dense)             (None, 1)                 33        \n",
      "=================================================================\n",
      "Total params: 10,561\n",
      "Trainable params: 10,561\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n",
      "Train on 3551 samples, validate on 888 samples\n",
      "Epoch 1/10\n",
      "3551/3551 [==============================] - 58s 16ms/step - loss: 0.3651 - acc: 0.8736 - val_loss: 0.1433 - val_acc: 0.9595\n",
      "Epoch 2/10\n",
      "3551/3551 [==============================] - 52s 15ms/step - loss: 0.1115 - acc: 0.9699 - val_loss: 0.0760 - val_acc: 0.9797\n",
      "Epoch 3/10\n",
      "3551/3551 [==============================] - 57s 16ms/step - loss: 0.0822 - acc: 0.9761 - val_loss: 0.0674 - val_acc: 0.9786\n",
      "Epoch 4/10\n",
      "3551/3551 [==============================] - 55s 16ms/step - loss: 0.0607 - acc: 0.9806 - val_loss: 0.0514 - val_acc: 0.9865\n",
      "Epoch 5/10\n",
      "3551/3551 [==============================] - 54s 15ms/step - loss: 0.0487 - acc: 0.9854 - val_loss: 0.0461 - val_acc: 0.9876\n",
      "Epoch 6/10\n",
      "3551/3551 [==============================] - 54s 15ms/step - loss: 0.0361 - acc: 0.9907 - val_loss: 0.0422 - val_acc: 0.9887\n",
      "Epoch 7/10\n",
      "3551/3551 [==============================] - 54s 15ms/step - loss: 0.0402 - acc: 0.9885 - val_loss: 0.0354 - val_acc: 0.9910\n",
      "Epoch 8/10\n",
      "3551/3551 [==============================] - 54s 15ms/step - loss: 0.0317 - acc: 0.9907 - val_loss: 0.0337 - val_acc: 0.9921\n",
      "Epoch 9/10\n",
      "3551/3551 [==============================] - 54s 15ms/step - loss: 0.0267 - acc: 0.9904 - val_loss: 0.0257 - val_acc: 0.9944\n",
      "Epoch 10/10\n",
      "3551/3551 [==============================] - 54s 15ms/step - loss: 0.0280 - acc: 0.9916 - val_loss: 0.0256 - val_acc: 0.9932\n",
      "[[0]\n",
      " [0]\n",
      " [0]]\n"
     ]
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "from __future__ import division, absolute_import\n",
    "\n",
    "import tldextract\n",
    "#import tflearn\n",
    "import os\n",
    "import pickle\n",
    "import numpy as np\n",
    "#from tflearn.data_utils import to_categorical, pad_sequences\n",
    "#import tensorflow as tf\n",
    "#import tflearn\n",
    "#from tflearn.layers.core import input_data, dropout, fully_connected\n",
    "#from tflearn.layers.conv import conv_1d, global_max_pool\n",
    "#from tflearn.layers.merge_ops import merge\n",
    "#from tflearn.layers.estimator import regression\n",
    "#from tflearn.data_utils import to_categorical, pad_sequences\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pickle\n",
    "from keras.utils import to_categorical\n",
    "from keras.models import Sequential,load_model\n",
    "from keras.layers import Dense, Embedding\n",
    "from keras.layers import LSTM,Bidirectional,GRU\n",
    "SRC_IP_IDX = 3-1               \n",
    "DST_IP_IDX = 4-1               \n",
    "SRC_PORT_IDX = 5-1             \n",
    "DST_PORT_IDX = 6-1             \n",
    "PROTOCOL_IDX = 7-1             \n",
    "DNS_QUERY_NAME_IDX = 55-1 # domain\n",
    "DNS_REQUEST_TYPE = 56-1\n",
    "DNS_DOMAIN_TTL = 59-1\n",
    "DNS_REPLY_IPV4IP = 60-1        \n",
    "DNS_REPLY_IPV6IP = 61-1        \n",
    "DNS_REPLY_RRTYPE = 62-1        \n",
    "DNS_REQUEST_LEN  = 88-1        \n",
    "DNS_REPLY_LENGTH = 90-1        \n",
    "  \n",
    "\n",
    "def iterbrowse(path):\n",
    "    for home, dirs, files in os.walk(path):\n",
    "        for filename in files:\n",
    "            yield os.path.join(home, filename)\n",
    "\n",
    "def pad_sequences(X, maxlen, value=0):\n",
    "    S=[]\n",
    "    for x in X:\n",
    "        xlen = len(x)\n",
    "        if xlen < maxlen:\n",
    "            x.extend([value]*(maxlen-xlen))\n",
    "        else:\n",
    "            x = x[:maxlen]\n",
    "        S.append(x)\n",
    "    return S\n",
    "def extract_domain(domain):\n",
    "    suffix = {'.com','.la','.io', '.co', '.cn','.info', '.net', '.org','.me', '.mobi', '.us', '.biz', '.xxx', '.ca', '.co.jp', '.com.cn', '.net.cn', '.org.cn', '.mx','.tv', '.ws', '.ag', '.com.ag', '.net.ag', '.org.ag','.am','.asia', '.at', '.be', '.com.br', '.net.br', '.name', '.live', '.news', '.bz', '.tech', '.pub', '.wang', '.space', '.top', '.xin', '.social', '.date', '.site', '.red', '.studio', '.link', '.online', '.help', '.kr', '.club', '.com.bz', '.net.bz', '.cc', '.band', '.market', '.com.co', '.net.co', '.nom.co', '.lawyer', '.de', '.es', '.com.es', '.nom.es', '.org.es', '.eu', '.wiki', '.design', '.software', '.fm', '.fr', '.gs', '.in', '.co.in', '.firm.in', '.gen.in', '.ind.in', '.net.in', '.org.in', '.it', '.jobs', '.jp', '.ms', '.com.mx', '.nl','.nu','.co.nz','.net.nz', '.org.nz', '.se', '.tc', '.tk', '.tw', '.com.tw', '.idv.tw', '.org.tw', '.hk', '.co.uk', '.me.uk', '.org.uk', '.vg'}\n",
    "\n",
    "    domain = domain.lower()\n",
    "    names = domain.split(\".\")\n",
    "    if len(names) >= 3:\n",
    "        if (\".\"+\".\".join(names[-2:])) in suffix:\n",
    "            return \".\".join(names[-3:]), \".\".join(names[:-3])\n",
    "        elif (\".\"+names[-1]) in suffix:\n",
    "            return \".\".join(names[-2:]), \".\".join(names[:-2])\n",
    "    #print (\"New domain suffix found. Use tld extract domain...\")\n",
    "\n",
    "    pos = domain.rfind(\"/\")\n",
    "    if pos >= 0: # maybe subdomain contains /, for dns tunnel tool\n",
    "        ext = tldextract.extract(domain[pos+1:])\n",
    "        subdomain = domain[:pos+1] + ext.subdomain\n",
    "    else:\n",
    "        ext = tldextract.extract(domain)\n",
    "        subdomain = ext.subdomain\n",
    "    if ext.suffix:\n",
    "        mdomain = ext.domain + \".\" + ext.suffix\n",
    "    else:\n",
    "        mdomain = ext.domain\n",
    "    return mdomain, subdomain\n",
    "\n",
    "\n",
    "def filter_metadata_dns(data):\n",
    "    if(len(data) < 91):\n",
    "        return False\n",
    "\n",
    "    protol  = data[PROTOCOL_IDX]\n",
    "    dstport = data[DST_PORT_IDX]\n",
    "    dstip   = data[DST_IP_IDX]\n",
    "    qname   = data[DNS_QUERY_NAME_IDX]\n",
    "\n",
    "    if '' == qname or '' == dstip:\n",
    "        return False\n",
    "    if '17' == protol and ('53' == dstport):\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "\n",
    "def metadata2_domain_data(log): \n",
    "    data = log.split('^')\n",
    "    if not filter_metadata_dns(data):\n",
    "        return None, None\n",
    "    domain = data[DNS_QUERY_NAME_IDX]\n",
    "    mdomain, subdomain = extract_domain(domain)\n",
    "    return (mdomain, subdomain)\n",
    "\n",
    "\n",
    "def get_local_data(tag=\"labeled\"):\n",
    "    data_path = \"./sample_data\"\n",
    "    black_data, cdn_data, white_data = [], [], []    \n",
    "    for dir_name in (\"black\", \"cdn\", \"white\"):\n",
    "        dir_path = \"%s/%s_%s\" % (data_path, tag, dir_name)\n",
    "\n",
    "        for path in iterbrowse(dir_path):\n",
    "            print(path)\n",
    "            with open(path) as f:\n",
    "                for line in f:\n",
    "                    mdomain, subdomain = metadata2_domain_data(line)\n",
    "                    if subdomain is not None:\n",
    "                        if \"white\" in path:\n",
    "                            white_data.append(subdomain)\n",
    "                        elif \"cdn\" in path:\n",
    "                            cdn_data.append(subdomain)\n",
    "                        elif \"black\" in path and \"pcap\" in path:\n",
    "                            black_data.append(subdomain)\n",
    "                        else:\n",
    "                            pass\n",
    "                            #print (\"pass path:\", path)\n",
    "                    #else:\n",
    "                    #    print (\"unknown line:\", line, \" in file:\", path)\n",
    "    return black_data, cdn_data, white_data\n",
    "\n",
    "\n",
    "class LABEL(object):\n",
    "    white = 0\n",
    "    cdn = 0\n",
    "    black = 1\n",
    "\n",
    "\n",
    "def get_data():\n",
    "    black_x, cdn_x, white_x = get_local_data()\n",
    "    black_y, cdn_y, white_y = [LABEL.black]*len(black_x), [LABEL.cdn]*len(cdn_x), [LABEL.white]*len(white_x)\n",
    "\n",
    "    X = black_x + cdn_x + white_x\n",
    "    labels = black_y + cdn_y + white_y\n",
    "\n",
    "    # Generate a dictionary of valid characters\n",
    "    valid_chars = {x:idx+1 for idx, x in enumerate(set(''.join(X)))}\n",
    "\n",
    "    max_features = len(valid_chars) + 1\n",
    "    #print \"max_features:\", max_features\n",
    "    maxlen = np.max([len(x) for x in X])\n",
    "    #print \"max_len:\", maxlen\n",
    "    maxlen = min(maxlen, 256)\n",
    "\n",
    "    # Convert characters to int and pad\n",
    "    X = [[valid_chars[y] for y in x] for x in X]\n",
    "    X = pad_sequences(X, maxlen=maxlen, value=0.)\n",
    "\n",
    "    # Convert labels to 0-1\n",
    "    #print(labels)\n",
    "    Y = labels\n",
    "    \n",
    "    volcab_file = \"volcab.pkl\"\n",
    "    output = open(volcab_file, 'wb')\n",
    "    # Pickle dictionary using protocol 0.\n",
    "    data = {\"valid_chars\": valid_chars, \"max_len\": maxlen, \"volcab_size\": max_features}\n",
    "    pickle.dump(data, output)\n",
    "    output.close()\n",
    "\n",
    "    return X, Y, maxlen, max_features\n",
    "def build_model_BiRNN(max_len, volcab_size):\n",
    "    \"\"\"Build Bi-RNN model\"\"\"\n",
    "    model = Sequential()\n",
    "    model.add(Embedding(input_dim=volcab_size,\n",
    "                        output_dim=64,\n",
    "                        input_length=max_len))\n",
    "    #model.add(LSTM(128, dropout=0.2, recurrent_dropout=0.2))\n",
    "    model.add(Bidirectional(GRU(16)))  \n",
    "    model.add(Dense(1, activation='sigmoid'))\n",
    "    #model.summary()\n",
    "    model.compile(loss='binary_crossentropy',\n",
    "                  optimizer='rmsprop',\n",
    "                  metrics=['accuracy'])\n",
    "    return model\n",
    "\n",
    "def get_rnn_model(max_len, volcab_size):\n",
    "    # Network building\n",
    "    net = tflearn.input_data([None, max_len])\n",
    "    net = tflearn.embedding(net, input_dim=volcab_size, output_dim=64)\n",
    "    net = tflearn.lstm(net, 64, dropout=0.8)\n",
    "    net = tflearn.fully_connected(net, 3, activation='softmax')\n",
    "    net = tflearn.regression(net, optimizer='adam', learning_rate=0.001,\n",
    "\t\t\t     loss='categorical_crossentropy')\n",
    "    model = tflearn.DNN(net, tensorboard_verbose=0)\n",
    "    return model\n",
    "\n",
    "\n",
    "def get_cnn_model(max_len, volcab_size):\n",
    "    # Building convolutional network\n",
    "    network = tflearn.input_data(shape=[None, max_len], name='input')\n",
    "    network = tflearn.embedding(network, input_dim=volcab_size, output_dim=64)\n",
    "    branch1 = conv_1d(network, 128, 3, padding='valid', activation='relu', regularizer=\"L2\")\n",
    "    branch2 = conv_1d(network, 128, 4, padding='valid', activation='relu', regularizer=\"L2\")\n",
    "    branch3 = conv_1d(network, 128, 5, padding='valid', activation='relu', regularizer=\"L2\")\n",
    "    network = merge([branch1, branch2, branch3], mode='concat', axis=1)\n",
    "    network = tf.expand_dims(network, 2)\n",
    "    network = global_max_pool(network)\n",
    "    network = dropout(network, 0.5)\n",
    "    network = fully_connected(network, 3, activation='softmax')\n",
    "    network = regression(network, optimizer='adam', learning_rate=0.001,\n",
    "                         loss='categorical_crossentropy', name='target')\n",
    "    model = tflearn.DNN(network, tensorboard_verbose=0)\n",
    "    return model\n",
    "\n",
    "\n",
    "def run():\n",
    "    X, Y, max_len, volcab_size = get_data()\n",
    "\n",
    "    print(\"X len:\", len(X), \"Y len:\", len(Y))\n",
    "    trainX, testX, trainY, testY = train_test_split(X, Y, test_size=0.2, random_state=42)\n",
    "    trainX=np.array(trainX)\n",
    "    testX=np.array(testX)\n",
    "    trainY=np.array(trainY)\n",
    "    testY=np.array(testY)\n",
    "    #print(trainX[:1])\n",
    "    #print(trainY[:1])\n",
    "    #print(testX[-1:])\n",
    "    #print(testY[-1:])\n",
    "    print(type(trainX))\n",
    "    model = build_model_BiRNN(max_len, volcab_size)\n",
    "    model.summary()\n",
    "    model.fit(trainX, trainY, validation_data=(testX, testY),batch_size=32,epochs=10)\n",
    "   \n",
    "    filename = 'dnstunnel2.module'\n",
    "    model.save(filename)\n",
    "\n",
    "    model=load_model(filename)\n",
    "    #print \"Just review 3 sample data test result:\"\n",
    "    result = model.predict_classes(testX[0:3])\n",
    "    print(result)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainX, testX, trainY, testY = train_test_split(X, Y, test_size=0.2, random_state=42)\n",
    "#print(type(trainX[0][0]))\n",
    "#print(len(trainX))\n",
    "#print(len(trainY))\n",
    "#print(type(trainY))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3551, 256)\n",
      "(3551,)\n"
     ]
    }
   ],
   "source": [
    "print(np.shape(trainX))\n",
    "print(np.shape(trainY))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3551, 256)\n",
      "(3551,)\n"
     ]
    }
   ],
   "source": [
    "trainX=np.array(trainX)\n",
    "#print((trainX[0:2]))\n",
    "#print(type(trainX))\n",
    "#print(type(trainX[0:2]))\n",
    "testX=np.array(testX)\n",
    "print(np.shape(trainX))\n",
    "trainY=np.array(trainY)\n",
    "#trainY=trainY.reshape(trainY.shape[0],1)\n",
    "#testY=np.array(testY)\n",
    "#testY=testY.reshape(testY.shape[0],1)\n",
    "#print((trainY))\n",
    "print(np.shape(trainY))\n",
    "#trainY=np.mat(trainY).flatten().T\n",
    "#testY=np.mat(testY).flatten().T\n",
    "#print(np.shape(trainY))\n",
    "#print(trainY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'numpy.ndarray'>\n",
      "<class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "print(type(trainX))\n",
    "print(type(trainY))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model_BiRNN(max_len, volcab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding_2 (Embedding)      (None, 256, 64)           2752      \n",
      "_________________________________________________________________\n",
      "lstm_2 (LSTM)                (None, 128)               98816     \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 1)                 129       \n",
      "=================================================================\n",
      "Total params: 101,697\n",
      "Trainable params: 101,697\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 3551 samples, validate on 888 samples\n",
      "Epoch 1/2\n",
      "3551/3551 [==============================] - 283s 80ms/step - loss: 0.6272 - acc: 0.6404 - val_loss: 0.6150 - val_acc: 0.6306\n",
      "Epoch 2/2\n",
      "3551/3551 [==============================] - 288s 81ms/step - loss: 0.5048 - acc: 0.7437 - val_loss: 0.4170 - val_acc: 0.8086\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f7c98cc3358>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(trainX, trainY,batch_size=32,epochs=2,validation_data=(testX, testY))\n",
    "#model.fit(trainX, trainY, validation_set=(testX, testY), show_metric=True, batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
